from django.contrib.auth import authenticate
from django.contrib.auth.models import User
from rest_framework import viewsets, status
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework.authtoken.models import Token
from rest_framework.permissions import IsAuthenticated
from rest_framework.decorators import api_view, permission_classes
from .models import Radiologist, Schedule, Leave, UnforeseenCircumstance, Study, Workload
from .serializers import (
    RadiologistSerializer, ScheduleSerializer, LeaveSerializer, UnforeseenCircumstanceSerializer, 
    StudySerializer, WorkloadSerializer
)
from .ml.main import predict_workload

class RadiologistViewSet(viewsets.ModelViewSet):
    queryset = Radiologist.objects.all()
    serializer_class = RadiologistSerializer
    # permission_classes = [permissions.IsAuthenticated]

class ScheduleViewSet(viewsets.ModelViewSet):
    queryset = Schedule.objects.all()
    serializer_class = ScheduleSerializer
    # permission_classes = [permissions.IsAuthenticated]

class LeaveViewSet(viewsets.ModelViewSet):
    queryset = Leave.objects.all()
    serializer_class = LeaveSerializer
    # permission_classes = [permissions.IsAuthenticated]

class UnforeseenCircumstanceViewSet(viewsets.ModelViewSet):
    queryset = UnforeseenCircumstance.objects.all()
    serializer_class = UnforeseenCircumstanceSerializer
    # permission_classes = [permissions.IsAuthenticated]

class StudyViewSet(viewsets.ModelViewSet):
    queryset = Study.objects.all()
    serializer_class = StudySerializer
    # permission_classes = [permissions.IsAuthenticated]

class WorkloadViewSet(viewsets.ModelViewSet):
    queryset = Workload.objects.all()
    serializer_class = WorkloadSerializer
    # permission_classes = [permissions.IsAuthenticated]

class PredictWorkloadAPI(APIView):
    def post(self, request):
        # Примерное тело запроса:
        # {'year': 2024, 'week_number': 25}
        year = request.data.get('year')
        week_number = request.data.get('week_number')
        # Прогнозирование нагрузки на основе предоставленных данных
        predicted_workload = predict_workload(year, week_number)
        # predicted_workload = {'modality': ['ct'], 'predicted_quantity': [1], 'predicted_ue': [1]}
        return Response(predicted_workload)
    

@api_view(['POST'])
def login(request):
    username = request.data.get('username')
    password = request.data.get('password')
    user = authenticate(username=username, password=password)
    if user is not None:
        token, _ = Token.objects.get_or_create(user=user)
        return Response({'token': token.key, 'user': {'username': user.username}}, status=status.HTTP_200_OK)
    return Response({'error': 'Invalid Credentials'}, status=status.HTTP_400_BAD_REQUEST)

@api_view(['POST'])
def register(request):
    username = request.data.get('username')
    password = request.data.get('password')
    if User.objects.filter(username=username).exists():
        return Response({'error': 'Username already exists'}, status=status.HTTP_400_BAD_REQUEST)
    user = User.objects.create_user(username=username, password=password)
    token, _ = Token.objects.get_or_create(user=user)
    return Response({'token': token.key, 'user': {'username': user.username}}, status=status.HTTP_201_CREATED)

@api_view(['GET'])
@permission_classes([IsAuthenticated])
def get_user(request):
    user = request.user
    return Response({
        'username': user.username,
        'email': user.email,
        'first_name': user.first_name,
        'last_name': user.last_name
    })